#include "flip_custom_tiling.h"
#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

class KernelFlip {
public:
    __aicore__ inline KernelFlip() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z, FlipCustomTilingData tiling)
    {
        this->dim = tiling.dim;
        this->dimN = tiling.dimN;
        this->dimC = tiling.dimC;
        this->dimH = tiling.dimH;
        this->dimW = tiling.dimW;
        this->blockLength = tiling.totalLength / AscendC::GetBlockNum();
        this->tileNum = tiling.tileNum;
        this->tileLength = this->blockLength / tiling.tileNum / BUFFER_NUM;
        xGm.SetGlobalBuffer((__gm__ float *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        zGm.SetGlobalBuffer((__gm__ float *)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(workBuffer, this->tileLength * sizeof(float));

    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress)
    {
        switch(this->dim) {
            case 0:
                ComputeDim0(progress);
                break;
            case 1:
                ComputeDim1(progress);
                break;
            case 2:
                ComputeDim2(progress);
                break;
            case 3:
                ComputeDim3(progress);
                break;
            default:
                return ;
                break;
        
        }    
    }

    __aicore__ inline void ComputeDim0(int32_t progress)
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();

        AscendC::DataCopy(zLocal, xLocal, this->tileLength);
        
        outQueueZ.EnQue<float>(zLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void ComputeDim1(int32_t progress)
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();

        AscendC::DataCopy(zLocal, xLocal, this->tileLength);
        
        outQueueZ.EnQue<float>(zLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void ComputeDim2(int32_t progress)
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();

        AscendC::DataCopy(zLocal, xLocal, this->tileLength);
        
        outQueueZ.EnQue<float>(zLocal);
        inQueueX.FreeTensor(xLocal);
    }
        
    __aicore__ inline void ComputeDim3(int32_t progress) 
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();
        
        // 分配偏移量张量
        AscendC::LocalTensor<uint32_t> offsetLocal = workBuffer.Get<uint32_t>();
        
        // 生成翻转的偏移量索引（以字节为单位）
        for (int32_t i = 0; i < this->tileLength; i++) {
            uint32_t reverseIndex = (this->tileLength - 1 - i) * sizeof(float);
            offsetLocal.SetValue(i, reverseIndex);
        }
        
        // 使用Gather算子进行向量化翻转
        AscendC::Gather(zLocal, xLocal, offsetLocal, (uint32_t)0, this->tileLength);
        
        // 清理资源
        outQueueZ.FreeTensor(offsetLocal);
        outQueueZ.EnQue<float>(zLocal);
        inQueueX.FreeTensor(xLocal);
    }

    // 使用循环实现元素翻转
    // __aicore__ inline void ComputeDim3(int32_t progress)
    // {
    //     AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
    //     AscendC::LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();

    //     for (int32_t i = 0; i < this->tileLength; i++) 
    //     {
    //         int32_t flippedPos = this->tileLength - 1 - i;
    //         auto xValue = xLocal.GetValue(flippedPos);
    //         zLocal.SetValue(i, xValue);
    //     }
            
    //     outQueueZ.EnQue<float>(zLocal);
    //     inQueueX.FreeTensor(xLocal);
    // }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<float> zLocal = outQueueZ.DeQue<float>();
        
        int32_t globalTileOffset = AscendC::GetBlockIdx() * this->blockLength + progress * this->tileLength;
        int32_t outputOffset = 0;
        
        switch (this->dim) {
            case 0: // 翻转N维度
                outputOffset = CalculateDim0OutputOffset(globalTileOffset, progress);
                break;
            case 1: // 翻转C维度
                outputOffset = CalculateDim1OutputOffset(globalTileOffset, progress);
                break;
            case 2: // 翻转H维度
                outputOffset = CalculateDim2OutputOffset(globalTileOffset, progress);
                break;
            case 3: // 翻转W维度
                outputOffset = CalculateDim3OutputOffset(globalTileOffset, progress);
                break;
            default:
                // 默认不翻转
                outputOffset = progress * this->tileLength;
                break;
        }
        
        // 计算在当前block内的偏移
        // int32_t blockOutputOffset = outputOffset - AscendC::GetBlockIdx() * this->blockLength;
        AscendC::DataCopy(zGm[outputOffset], zLocal, this->tileLength);
        
        outQueueZ.FreeTensor(zLocal);
    }

    // 计算N维度翻转的输出偏移量
    __aicore__ inline int32_t CalculateDim0OutputOffset(int32_t globalOffset, int32_t progress)
    {
        int32_t absolute_tile_index = AscendC::GetBlockIdx() * this->tileNum * BUFFER_NUM + progress;
        int32_t global_start_pos = absolute_tile_index * this->tileLength;
        
        // 计算原始NCHW坐标
        int32_t n_index = global_start_pos / (this->dimC * this->dimH * this->dimW);
        int32_t chw_offset = global_start_pos % (this->dimC * this->dimH * this->dimW);
        
        // N维度翻转
        int32_t flipped_n = this->dimN - 1 - n_index;
        
        // 计算翻转后的全局位置
        int32_t flipped_global = flipped_n * (this->dimC * this->dimH * this->dimW) + chw_offset;
        
        return flipped_global - AscendC::GetBlockIdx() * this->blockLength;
    }

    // 计算C维度翻转的输出偏移量
    __aicore__ inline int32_t CalculateDim1OutputOffset(int32_t globalOffset, int32_t progress)
    {
        int32_t absolute_tile_index = AscendC::GetBlockIdx() * this->tileNum * BUFFER_NUM + progress;
        int32_t global_start_pos = absolute_tile_index * this->tileLength;
        
        // 计算原始NCHW坐标
        int32_t n_index = global_start_pos / (this->dimC * this->dimH * this->dimW);
        int32_t remaining = global_start_pos % (this->dimC * this->dimH * this->dimW);
        int32_t c_index = remaining / (this->dimH * this->dimW);
        int32_t hw_offset = remaining % (this->dimH * this->dimW);
        
        // C维度翻转
        int32_t flipped_c = this->dimC - 1 - c_index;
        
        // 计算翻转后的全局位置
        int32_t flipped_global = n_index * (this->dimC * this->dimH * this->dimW) + flipped_c * (this->dimH * this->dimW) + hw_offset;
        
        return flipped_global - AscendC::GetBlockIdx() * this->blockLength;
    }

    // 计算H维度翻转的输出偏移量
    __aicore__ inline int32_t CalculateDim2OutputOffset(int32_t globalOffset, int32_t progress)
    {
        int32_t absolute_tile_index = AscendC::GetBlockIdx() * this->tileNum * BUFFER_NUM + progress;
        int32_t global_start_pos = absolute_tile_index * this->tileLength;
        
        // 计算原始NCHW坐标
        int32_t n_index = global_start_pos / (this->dimC * this->dimH * this->dimW);
        int32_t remaining = global_start_pos % (this->dimC * this->dimH * this->dimW);
        int32_t c_index = remaining / (this->dimH * this->dimW);
        remaining = remaining % (this->dimH * this->dimW);
        int32_t h_index = remaining / this->dimW;
        int32_t w_offset = remaining % this->dimW;
        
        // H维度翻转
        int32_t flipped_h = this->dimH - 1 - h_index;
        
        // 计算翻转后的全局位置
        int32_t flipped_global = n_index * (this->dimC * this->dimH * this->dimW) + c_index * (this->dimH * this->dimW) + 
                            flipped_h * this->dimW + w_offset;
        
        return flipped_global - AscendC::GetBlockIdx() * this->blockLength;
    }

    // 计算W维度翻转的输出偏移量
    __aicore__ inline int32_t CalculateDim3OutputOffset(int32_t globalOffset, int32_t progress)
    {
        int32_t absolute_tile_index = AscendC::GetBlockIdx() * this->tileNum * BUFFER_NUM + progress;
        int32_t global_start_pos = absolute_tile_index * this->tileLength;
        
        // 计算原始NCHW坐标
        int32_t n_index = global_start_pos / (this->dimC * this->dimH * this->dimW);
        int32_t remaining = global_start_pos % (this->dimC * this->dimH * this->dimW);
        int32_t c_index = remaining / (this->dimH * this->dimW);
        remaining = remaining % (this->dimH * this->dimW);
        int32_t h_index = remaining / this->dimW;
        int32_t w_start = remaining % this->dimW;
        
        // 计算tile在W维度内的位置
        int32_t tiles_per_row = this->dimW / this->tileLength;
        int32_t tile_index_in_w = w_start / this->tileLength;
        
        // W维度tile间翻转
        int32_t flipped_tile_index = tiles_per_row - 1 - tile_index_in_w;
        
        // 计算翻转后的全局位置
        int32_t flipped_global = n_index * (this->dimC * this->dimH * this->dimW) + c_index * (this->dimH * this->dimW) + 
                            h_index * this->dimW + flipped_tile_index * this->tileLength;
        
        return flipped_global - AscendC::GetBlockIdx() * this->blockLength;
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> workBuffer;
    AscendC::GlobalTensor<float> xGm;
    AscendC::GlobalTensor<float> zGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    int32_t dim;
    uint32_t dimN;
    uint32_t dimC;
    uint32_t dimH;
    uint32_t dimW;
};

extern "C" __global__ __aicore__ void flip_custom(GM_ADDR x, GM_ADDR z, FlipCustomTilingData tiling)
{
    KernelFlip op;
    op.Init(x, z, tiling);
    op.Process();
}
